# Simple types and type checking in Python.


class Exp:
    pass


class Lam(Exp):
    def __init__(self, vartype, body):
        self.body = body
        self.vartype = vartype
        self.varname = name = self.body.func_code.co_varnames[0]

    def __repr__(self):
        name = self.varname
        return "Lam(%r, lambda %s: %r)" % (self.vartype, name, self.body(Var(name)))

    def __str__(self, paren=False):
        name = self.varname
        s = "lam %s:%s. %s" % (name, self.vartype,
                               self.body(Var(name)).__str__(False))
        if paren:
            s = "(" + s + ")"
        return s

    def __call__(self, arg):
        return self.body(arg)

    def eval_(self):
        return self

    def check(self, type_):
        if not isinstance(type_, Arrow):
            return False
        left, right = type_.left, type_.right
        name = self.varname
        # We pass the body a variable object with the appropriate type
        # and type check the result.
        return self.vartype == left and self.body(Var(name, left)).check(right)
    

class App(Exp):
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def __repr__(self):
        return "App(%r, %r)" % (self.left, self.right)

    def __str__(self, paren=False):
        return "%s %s" % (self.left.__str__(True), self.right.__str__(True))

    def eval_(self):
        left = self.left.eval_()
        return left(self.right).eval_()

    def check(self, type_):
        left = self.left.vartype
        return self.left.check(Arrow(left, type_)) and self.right.check(left)        


class Var(Exp):
    """
    A dummy variable object used to analyze lambda bodies.
    Also contains the type of the variable.
    """
    def __init__(self, name, type_=None):
        self.name = name
        self.type = type_

    def __repr__(self):
        return self.name

    def __str__(self, paren):
        return self.name

    def check(self, type_):
        return (self.type == type_)


class Type:
    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return "Type(%r)" % self.name

    def __str__(self, paren=False):
        return self.name

    def __eq__(self, other):
        return self.name == other.name


class Arrow(Type):
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def __repr__(self):
        return "Arrow(%r, %r)" % (self.left, self.right)

    def __str__(self, paren=False):
        s = "%s->%s" % (self.left.__str__(True), self.right)
        if paren:
            s = "(" + s + ")"
        return s

    def __eq__(self, other):
        if not isinstance(other, Arrow):
            return False
        return self.left == other.left and self.right == other.right


# Examples

def check(exp, type_):
    if exp.check(type_):
        print "%s : %s" % (exp, type_)
    else:
        print "Type checking error."

a = Type("a")
I = Lam(a, lambda x: x)
K = Lam(a, lambda x: Lam(a, lambda y: x))
O = Lam(a, lambda x : App(x, x))

print "I =", I
print "K =", K
print "O =", O

check(I, Arrow(a, a))
check(Lam(Arrow(a, a), lambda x: x), Arrow(Arrow(a, a), Arrow(a, a)))
check(App(Lam(Arrow(a, a), lambda x: x), I), Arrow(a, a))
